# coding: utf-8
import argparse
import logging
import os
import data_task
from utils import gps_dict


"""
gen_origin_index:
生成原始数据的index
example:
python src/data.py gen_origin_index --origin-dir data/shanghai/L17/ --origin-index data/shanghai/L17.txt

gen_folded_img:
做数据集的增大.存储的文件夹不存在会自动创建!
example:
python src/data.py gen_folded_img --origin-index data/shanghai/L17.txt --folded-index data/shanghai/L17_append.txt --folded-dir data/shanghai/L17_append/ --folded-num 3 --region shanghai

test_map_build:
测试map，会对map做详细的检查。可以带fold。可以不带fold。带fold时指定index即可
example:
(no folding)
python src/data.py test_map_build --origin-index data/shanghai/L17.txt  --region shanghai
(with folding)  
python src/data.py test_map_build --origin-index data/shanghai/L17.txt --folded-index data/shanghai/L17_append.txt --folded-num 3 --region shanghai

test_pos_img:
测试某一点位置处的图片。
example:
python src/data.py test_pos_img --origin-index data/shanghai/L17.txt --folded-index data/shanghai/L17_append.txt --folded-num 3 --region shanghai --test-la 31.3016102468 --test-lo 121.5009278849
python src/data.py test_pos_img --origin-index data/shanghai/L17.txt --region shanghai --test-la 31.3016102468 --test-lo 121.5009278849

gen_mxnet_index:
创建mxnet的index
example:
python src/data.py gen_mxnet_index --origin-index data/shanghai/L17.txt --folded-index data/shanghai/L17_append.txt --folded-num 3 --region shanghai --y-index data/restaurant/Shanghai/shanghai_wgs.txt --mxnet-index data/mxnet_shanghai_L17_fold_3.txt --y-scale 0,1,2,10
用--balance-factor和 --balance-label控制某一类下采样比例
用--y-scale控制分类的间隔

gen_tensorflow_imgs:
创建tensorflow的那种图片，全程为复制
example:
python src/data.py gen_tensorflow_imgs --origin-index data/shanghai/L17.txt --folded-index data/shanghai/L17_append.txt --folded-num 3 --region shanghai --y-index data/restaurant/Shanghai/shanghai_wgs.txt --tensorflow-dir data/tensorflow_images/ --y-scale 0,1,5,10,20,30,60 --balance-factor 0.05 --tensorflow-train-factor 0.8
用--y-scale控制分类的间隔
用--balance-factor和 --balance-label控制某一类下采样比例
用--tensorflow-train-factor控制训练集百分比

gen_tensorflow_data:
根据上面tensorflow的文件夹创建tensorflow record
example:
python src/data.py gen_tensorflow_data --tensorflow-dir data/tensorflow_images/ 
用--tensorflow-train-shards和--tensorflow-validation-shards控制分片个数

"""


def arg_parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('task', choices=['gen_origin_index', 'gen_folded_img', 'test_map_build',
                                         'test_pos_img', 'gen_mxnet_index', 'gen_tensorflow_imgs', 'gen_tensorflow_data'])
    parser.add_argument('--origin-dir', help='file dir of original images. Like Shanghai/L17')
    parser.add_argument('--folded-dir', help='file dir of folded images. Like Shanghai/L17_append')
    parser.add_argument('--origin-index', help='index file for original images. Like Shanghai/L17.txt')
    parser.add_argument('--folded-index', help='index file for folded images. Like Shanghai/L17_append.txt')
    parser.add_argument('--folded-num', type=int, help='folded num')
    parser.add_argument('--left-up-la', help='left up corner\'s la', type=float)
    parser.add_argument('--left-up-lo', help='left up corner\'s lo', type=float)
    parser.add_argument('--right-down-la', help='right down corner\'s la', type=float)
    parser.add_argument('--right-down-lo', help='right down corner\'s lo', type=float)
    parser.add_argument('--test-la', help='test point\'s la(wgs84)', type=float, default=None)
    parser.add_argument('--test-lo', help='test point\'s lo(wgs84)', type=float, default=None)

    parser.add_argument('--mxnet-index', help='index file for mxnet. Like Shanghai/mxnet_L17.txt')
    parser.add_argument('--y-index', help='y index(only WGS84). Like data/restaurant/shanghai/shanghai_wgs.txt')
    parser.add_argument('--y-scale', help='y scale, default to 0,1. Example: scale [a,b,c,d] 1. >=a <b 2. >=b <c 3. >=c <d 4. >=d', default='0,1')
    parser.add_argument('--tensorflow-dir', help='dir for tensorflow. Like Shanghai/TF_format_L17/')
    parser.add_argument('--tensorflow-train-factor', help='train sample factor for tensorflow.default=1', type=float, default=1.0)
    parser.add_argument('--tensorflow-train-shards', help='', type=int, default=2)
    parser.add_argument('--tensorflow-validation-shards', help='', type=int, default=2)

    parser.add_argument('--balance-factor', help='balance_factor', type=float, default='1.0')
    parser.add_argument('--balance-label', help=' balance_label', type=float, default='0')

    # region
    parser.add_argument('--region', help='short cut for those la and lo', type=str, choices=['shanghai', 'beijing'])

    return parser.parse_args()


def logging_config():
    # logging.basicConfig(level=logging.INFO, format='[%(levelname)s]%(module)s->%(funcName)s: %(message)s --- %(asctime)s')
    try:
        import colorlog
        have_colorlog = True
    except ImportError:
        have_colorlog = False

    log = logging.getLogger()
    log.setLevel(logging.INFO)
    format = '[%(levelname)s]%(module)s->%(funcName)s: %(message)s --- %(asctime)s'
    cformat = '%(log_color)s' + format
    date_format = '%Y-%m-%d %H:%M:%S'
    f = logging.Formatter(format, date_format)
    if have_colorlog is True:
        cf = colorlog.ColoredFormatter(cformat, date_format,
                                       log_colors={'DEBUG': 'reset', 'INFO': 'reset',
                                                   'WARNING': 'bold_yellow', 'ERROR': 'bold_red',
                                                   'CRITICAL': 'bold_red'})
    else:
        cf = f
    ch = logging.StreamHandler()
    if os.isatty(2):
        ch.setFormatter(cf)
    else:
        ch.setFormatter(f)
    log.addHandler(ch)


def main():
    args = arg_parse()
    logging_config()

    # --> abs path
    current_path = os.getcwd()
    if args.origin_dir is not None and os.path.isabs(args.origin_dir) is False:
        args.origin_dir = os.path.join(current_path, args.origin_dir)
    if args.folded_dir is not None and os.path.isabs(args.folded_dir) is False:
        args.folded_dir = os.path.join(current_path, args.folded_dir)

    # region -> wgs84
    if args.region is not None:
        assert args.left_up_la is None
        assert args.left_up_lo is None
        assert args.right_down_la is None
        assert args.right_down_lo is None
        args.left_up_la = gps_dict.data[args.region]['left_up_la']
        args.left_up_lo = gps_dict.data[args.region]['left_up_lo']
        args.right_down_la = gps_dict.data[args.region]['right_down_la']
        args.right_down_lo = gps_dict.data[args.region]['right_down_lo']

    # y-scale --> number list
    args.y_scale = args.y_scale.split(',')
    for i in range(len(args.y_scale)):
        assert args.y_scale[i] != ''
        args.y_scale[i] = int(args.y_scale[i])

    work = getattr(data_task, args.task)
    work(args)


if __name__ == '__main__':
    main()
